module Hexena.PEHeader 
  ( PEOptionalHeader (..)
  , PEHeader (..)
  , peheader2KeyValList
  , peoheader2KeyValList
  , readPE
  , readPEFrom
  , readPEIO
  ) where

import Control.Monad
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Word

import Hexena.MZHeader

data PEOptionalHeader = PEOptionalHeader {
      peohMagic :: Word16,
      peohMajorLinkerVersion :: Word8,
      peohMinorLinkerVersion :: Word8,
      peohSizeOfCode :: Word32,
      peohSizeofInitialisedData :: Word32,
      peohSizeofUninitialisedData :: Word32,
      peohAddressOfEntryPoint :: Word32,
      peohBaseOfCode :: Word32,
      peohBaseOfData :: Word32,
      peohImageBase :: Word32,
      peohSectionAlignment :: Word32,
      peohFileAlignment ::Word32,
      peohMajorOSVersion :: Word16,
      peohMinorOSVersion :: Word16,
      peohMajorImageVersion :: Word16,
      peohMinorImageVersion :: Word16,
      peohMajorSubsystemVersion :: Word16,
      peohMinorSubsystemVersion :: Word16,
      peohReserved :: Word32,
      peohSizeOfImage :: Word32,
      peohSizeOfHeaders :: Word32,
      peohCheckSum :: Word32
    } deriving (Show)

data PEHeader = PEHeader {
      peSignature :: Word32,
      peMachine :: Word16,
      peNumberOfSections :: Word16,
      peTimeDateStamp :: Word32,
      pePointerToSymbolTable :: Word32,
      peNumberOfSymbols :: Word32,
      peSizeOfOptionalHeader :: Word16,
      peCharacteristics :: Word16,
      peOptionalHeader :: PEOptionalHeader
    } deriving (Show)


readPEOH :: Get PEOptionalHeader
readPEOH = return PEOptionalHeader
           `ap` getWord16le
           `ap` getWord8 `ap` getWord8
           `ap` getWord32le `ap` getWord32le `ap` getWord32le
           `ap` getWord32le
           `ap` getWord32le `ap` getWord32le
           `ap` getWord32le
           `ap` getWord32le `ap` getWord32le -- section/file alignment
           `ap` getWord16le `ap` getWord16le `ap` getWord16le `ap` getWord16le `ap` getWord16le `ap` getWord16le
           `ap` getWord32le -- Reserved
           `ap` getWord32le `ap` getWord32le
           `ap` getWord32le -- CheckSum


-- | Read the PE Header from a file.
--
-- This will read and interpret the MZ header and try to find the
-- appropriate beginning of a PE header using lfanew in the MZ
-- header. If this fails (header is corrupt) it will search for the
-- first occurence of "PE" in the file and use this as a header. We
-- have a slight chance that this may work. BTW Windows happily
-- executes all this stuff...
readPEFrom :: BL.ByteString -> PEHeader
readPEFrom dat = let (_, lfanew) = runGet checkFile dat
                     bs x = let x' = BL.dropWhile (/= 0x50) x
                                x'' = BL.tail x'
                            in if 0x45 == BL.head x'' then x' else bs x''
                     drp x | x < BL.length dat = BL.drop x dat
                           | otherwise = bs dat
                 in runGet readPE $ drp $ fromIntegral lfanew


readPEIO :: String -> IO PEHeader
readPEIO x = readPEFrom `fmap` BL.readFile x

readPE :: Get PEHeader
readPE = do sig <- getWord32le
            mac <- getWord16le
            nos <- getWord16le
            tds <- getWord32le
            ptst <- getWord32le
            nosy <- getWord32le
            sooh <- getWord16le
            cha <- getWord16le
            peoh <- readPEOH
            return $ PEHeader sig mac nos tds ptst nosy sooh cha peoh
  
peheader2KeyValList :: PEHeader -> [(String, Int)]
peheader2KeyValList (PEHeader sig mac nos tds ptst nosy sooh cha _) = [
  ("Signature", fromIntegral sig),
  ("Machine", fromIntegral mac),
  ("NumberOfSections", fromIntegral nos),
  ("TimeDateStamp", fromIntegral tds),
  ("PointerToSymbolTable", fromIntegral ptst),
  ("NumberOfSymbols", fromIntegral nosy),
  ("SizeOfOptionalHeader", fromIntegral sooh),
  ("Characteristics", fromIntegral cha)
 ]

peoheader2KeyValList :: PEOptionalHeader -> [(String, Int)]
peoheader2KeyValList peoh = [
  ("Magic", fromIntegral $ peohMagic peoh),
  ("MajorLinkerVersion", fromIntegral $ peohMajorLinkerVersion peoh),
  ("MinorLinkerVersion", fromIntegral $ peohMinorLinkerVersion peoh),
  ("SizeOfCode", fromIntegral $ peohSizeOfCode peoh),
  ("SizeofInitialisedData", fromIntegral $ peohSizeofInitialisedData peoh),
  ("SizeofUninitialisedData", fromIntegral $ peohSizeofUninitialisedData peoh),
  ("AddressOfEntryPoint", fromIntegral $ peohAddressOfEntryPoint peoh),
  ("BaseOfCode", fromIntegral $ peohBaseOfCode peoh),
  ("BaseOfData", fromIntegral $ peohBaseOfData peoh),
  ("ImageBase", fromIntegral $ peohImageBase peoh),
  ("SectionAlignment", fromIntegral $ peohSectionAlignment peoh),
  ("FileAlignment", fromIntegral $ peohFileAlignment peoh),
  ("MajorOSVersion", fromIntegral $ peohMajorOSVersion peoh),
  ("MinorOSVersion", fromIntegral $ peohMinorOSVersion peoh),
  ("MajorImageVersion", fromIntegral $ peohMajorImageVersion peoh),
  ("MinorImageVersion", fromIntegral $ peohMinorImageVersion peoh),
  ("MajorSubsystemVersion", fromIntegral $ peohMajorSubsystemVersion peoh),
  ("MinorSubsystemVersion", fromIntegral $ peohMinorSubsystemVersion peoh),
  ("Reserved", fromIntegral $ peohReserved peoh),
  ("SizeOfImage", fromIntegral $ peohSizeOfImage peoh),
  ("SizeOfHeaders", fromIntegral $ peohSizeOfHeaders peoh),
  ("CheckSum", fromIntegral $ peohCheckSum peoh)
 ]